#!/usr/bin/python

# make the clustering class-specific

import sys,os,re,glob,math,glob,signal,traceback
import matplotlib
if "DISPLAY" not in os.environ: matplotlib.use("AGG")
else: matplotlib.use("GTK")
from scipy.ndimage import interpolation
from pylab import *
from optparse import OptionParser
from multiprocessing import Pool
import ocrolib
from ocrolib import number_of_processors,die

signal.signal(signal.SIGINT,lambda *args:sys.exit(1))

parser = OptionParser("""
usage: %prog [options] image1.png image2.png ...

Computes recognition lattices for text lines.  Also displays the bestpath
result (recognition result without language model).
""")

# these options control alignment
parser.add_option("-m","--model",help="model file",default="default.cmodel")
parser.add_option("-w","--whitespace",help="space model file",default="space.model")
parser.add_option("-c","--cont",help="continue even when unexpected errors occur",action="store_true")
parser.add_option("-R","--maxrange",help="max range for grouping",type=int,default=4)
parser.add_option("-L","--noligatures",help="don't output ligatures",action="store_true")
parser.add_option("-N","--nbest",help="use the nbest results from classifier",type=int,default=10)
parser.add_option("-C","--maxcost",help="maximum cost to be considered",type=float,default=10.0)
parser.add_option("-s","--segmenter",help="segmenter",default="DpSegmenter")
parser.add_option("-S","--suffix",help="suffix for writing rseg/cseg files",default=None)
parser.add_option("-O","--overwrite",help="overwrite rseg/cseg files",action="store_true")
parser.add_option("-Q","--parallel",type=int,default=number_of_processors(),help="number of parallel processes to use")
parser.add_option("-d","--display",help="display progress",action="store_true")
parser.add_option("-v","--verbose",help="verbose recognition output",action="store_true")
parser.add_option("-V","--verbose2",help="extra verbose recognition output",action="store_true")
parser.add_option("-Y","--examine",help="examine recognition for these classes",default="")
(options,args) = parser.parse_args()

if len(args)==0:
    parser.print_help()
    sys.exit(0)

if len(args)==1 and os.path.isdir(args[0]):
    args = sorted(glob.glob(args[0]+"/????/??????.png"))

if options.display:
    ion()

if options.verbose2: verbose = 2
elif options.verbose: verbose=1
else: verbose = 0

cmodel = ocrolib.load_component(ocrolib.ocropus_find_file(options.model))

assert cmodel is not None
if hasattr(cmodel,'recognizeLine'):
    linerec = cmodel
else:
    if hasattr(cmodel,'coutputs'):
        linerec = ocrolib.CmodelLineRecognizer(cmodel=cmodel,whitespace=options.whitespace,
                                               maxrange=options.maxrange,verbose=verbose,
                                               debug_cls=options.examine)
        linerec.use_ligatures = not options.noligatures
        linerec.maxcost = options.maxcost
        linerec.nbest = options.nbest
    else:
        print "%s (%s): don't know what kind of model this is"%(options.model,cmodel)

def extract(t):
    fname = t
    image = ocrolib.read_image_gray(fname)
    ocrolib.fcleanup(fname,options.suffix,["cseg","txt","costs","fst"])

    if options.display:
        clf()
        gray()
        imshow(image)
        ginput(1,timeout=1)
    try:
        lattice,rseg = linerec.recognizeLineSeg(image)
    except ocrolib.RecognitionError,e:
        print fname,"*** %s ***"%e
        return
    except:
        traceback.print_exc()
        if not options.cont: die("internal error")
        return

    lattice_file = ocrolib.fvariant(fname,"fst",options.suffix)
    if os.path.exists(lattice_file) and not options.overwrite:
        die("%s: %s",lattice_file,"already exists")
    lattice.save(lattice_file)

    rseg_file = ocrolib.fvariant(fname,"rseg",options.suffix)
    if os.path.exists(rseg_file) and not options.overwrite:
        die("%s: %s",rseg_file,"already exists")
    ocrolib.write_line_segmentation(rseg_file,rseg)

    s = lattice.bestpath()
    try:
        print fname,":",s
    except:
        print fname,":",repr(s)

if options.parallel<2:
    for arg in args: extract(arg)
else:
    pool = Pool(processes=options.parallel)
    result = pool.map(extract,args)
